PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction#41048
Open
copybara-service[bot] wants to merge 1 commit intomainfrom
Open
PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction#41048copybara-service[bot] wants to merge 1 commit intomainfrom
copybara-service[bot] wants to merge 1 commit intomainfrom
Conversation
bfcc606 to
acd48c2
Compare
…ll instruction Imported from GitHub PR #40920 - New `PropagateCallMetadata` HLO pass that propagates metadata (op_name prefix and stack_frame_id) from `kCall` instructions into their called computations, recursing through nested control-flow (while, conditional) but not into embedded computations (reduce's to_apply, etc.) - Removes metadata propagation from `CallInliner` — the `RecursivelyUpdateMetadata` helper and per-instruction metadata update during inlining are no longer needed since the new pass handles this as a separate concern - Wired into GPU compiler pre-SPMD pipeline before `CallInliner`, so metadata is propagated while `kCall` ops still exist (including non-inlinable calls) - Test coverage for op_name propagation, stack frame concatenation, overflow protection, redundant prefix detection, nested calls, and idempotency #### Motivation The `CallInliner` only updates metadata for calls it actually inlines. Non-inlinable calls (e.g. calls with inlineable="false") were skipped entirely, leaving their callee instructions with incomplete metadata context. Extracting this into a standalone pass ensures all calls get metadata propagation regardless of inlining decisions. Copybara import of the project: -- 3315b79 by Eugene Zhulenev <ezhulenev@openxla.org>: [xla] Add a separate pass to propagate metadata across kCall instructions Merging this change closes #40920 FUTURE_COPYBARA_INTEGRATE_REVIEW=#40920 from ezhulenev:propagate-metadata-pass 3315b79 PiperOrigin-RevId: 900707821
acd48c2 to
e0dbc47
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR #40920: [xla] Add a separate pass to propagate metadata across kCall instruction
Imported from GitHub PR #40920
PropagateCallMetadataHLO pass that propagates metadata (op_name prefix and stack_frame_id) fromkCallinstructions into their called computations, recursing through nested control-flow (while, conditional) but not into embedded computations (reduce's to_apply, etc.)CallInliner— theRecursivelyUpdateMetadatahelper and per-instruction metadata update during inlining are no longer needed since the new pass handles this as a separate concernCallInliner, so metadata is propagated whilekCallops still exist (including non-inlinable calls)Motivation
The
CallInlineronly updates metadata for calls it actually inlines. Non-inlinable calls (e.g. calls with inlineable="false") were skipped entirely, leaving their callee instructions with incomplete metadata context. Extracting this into a standalone pass ensures all calls get metadata propagation regardless of inlining decisions.Copybara import of the project:
--
3315b79 by Eugene Zhulenev ezhulenev@openxla.org:
[xla] Add a separate pass to propagate metadata across kCall instructions
Merging this change closes #40920
FUTURE_COPYBARA_INTEGRATE_REVIEW=#40920 from ezhulenev:propagate-metadata-pass 3315b79